from typing import Tuple, Dict, List
from omegaconf import DictConfig

import torch
import numpy as np
import pandas as pd

from torch.utils.data import Dataset

from ecgcmr.signal.sig_augmentations.ECGAugmentations import ECGAugmentations


class DownstreamECGDataset(Dataset):
  """
  Dataset for downstream task for ECG 
  """
  def __init__(
    self, 
    cfg: DictConfig, 
    mode: str,
    apply_augmentations: bool = False,
    ) -> None:
    
    task = cfg.downstream_task.task
    metrics_type = cfg.downstream_task.type

    if metrics_type == 'vol':
      num_classes = cfg.downstream_task.num_classes_vol
    elif metrics_type == 'area':
      num_classes = cfg.downstream_task.num_classes_area

    data_path = cfg.dataset.paths[f'data_ecg_{mode}']
    subset_idx_path = cfg.downstream_task.paths[f'downstream_task_{mode}_subindices']
    labels_path = cfg.downstream_task.paths[f'downstream_task_{mode}_labels_{metrics_type}']

    self.ecg_data = np.load(data_path, mmap_mode='c')
    self.labels = np.load(labels_path, mmap_mode='c')
    self.subset_idx = np.load(subset_idx_path, mmap_mode='c')
  
    if task == 'regression':
        self.labels = self.labels[:, :num_classes]

    self.ecg_augmentations = ECGAugmentations(
      cfg=cfg.downstream_task.augmentations.ecg,
      sampling_rate=cfg.downstream_task.sampling_rate,
      input_electrodes=cfg.downstream_task.input_electrodes,
      single_modality=cfg.downstream_task.single_modality,
      apply_augmentations=apply_augmentations,
      seed=cfg.seed
    )
    
  def __len__(self):
    return len(self.subset_idx)

  def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
    return self.ecg_augmentations(self.ecg_data[self.subset_idx[idx]]), torch.from_numpy(self.labels[idx]).float()



class DownstreamECGDatasetEvaluation(Dataset):
    """
    Dataset for downstream task for ECG with disease table integration.
    """
    def __init__(
        self, 
        cfg: DictConfig, 
        mode: str,
        disease_map: Dict[str, List[str]],
        apply_augmentations: bool = False,
        use_peaks_location: bool = True,
    ) -> None:
        task = cfg.downstream_task.task
        metrics_type = cfg.downstream_task.type

        if metrics_type == 'vol':
            num_classes = cfg.downstream_task.num_classes_vol
        elif metrics_type == 'area':
            num_classes = cfg.downstream_task.num_classes_area

        data_path = cfg.dataset.paths[f'data_ecg_{mode}']
        subset_idx_path = cfg.downstream_task.paths[f'downstream_task_{mode}_subindices']
        labels_path = cfg.downstream_task.paths[f'downstream_task_{mode}_labels_{metrics_type}']

        disease_table_path = cfg.downstream_task.paths[f'downstream_task_{mode}_disease_table']

        self.ecg_data = np.load(data_path, mmap_mode='c')
        self.labels = np.load(labels_path, mmap_mode='c')
        self.subset_idx = np.load(subset_idx_path, mmap_mode='c')

        # =========== Load disease table & combine codes ===========
        disease_table = pd.read_csv(disease_table_path)
        disease_table = disease_table.set_index("eid")
        disease_sub = disease_table.iloc[self.subset_idx].copy()  # shape [M, many_columns]

        self.diseases_of_interest = list(disease_map.keys())  # e.g. ["MI","CM","AF","HF"]

        for disease_name, code_list in disease_map.items():
            vals = np.zeros(len(disease_sub), dtype=np.float32)
            for c in code_list:
                if c not in disease_sub.columns:
                    raise KeyError(f"Column {c} not found in disease table!")
                col_vals = disease_sub[c].values  # shape [M,]
                vals = np.logical_or(vals, (col_vals==1))
            disease_sub[disease_name] = vals.astype(np.float32)

        final_df = disease_sub[self.diseases_of_interest]  # shape [M, len(disease_map)]
        self.disease_labels = final_df.values.astype(np.float32)   # [M, len(disease_map)]

        if task == 'regression':
            self.labels = self.labels[:, :num_classes]
            
        self.rpeaks_data = None
        paths = cfg.dataset.paths
        self.use_peaks_location = use_peaks_location
        
        if self.use_peaks_location:
            self.rpeaks_data = np.load(getattr(paths, f'data_ecg_rpeaks_updated_{mode}'), allow_pickle=True)
        else:
            raise ValueError("R-peaks path not found in cfg for 'use_peaks_location=True'.")

        # Augmentations
        self.ecg_augmentations = ECGAugmentations(
            cfg=cfg.downstream_task.augmentations.ecg,
            sampling_rate=cfg.downstream_task.sampling_rate,
            input_electrodes=cfg.downstream_task.input_electrodes,
            single_modality=False,
            use_peaks_location=self.use_peaks_location,
            apply_augmentations=apply_augmentations,
            seed=cfg.seed
        )

    def __len__(self):
        return len(self.subset_idx)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Returns:
            ecg_data_tensor : shape [C, T] or similar
            task_labels_t   : shape [num_classes], float
            disease_labels_t: shape [num_diseases], float(0/1)
            rpeaks_t        : 1D variable-length if use_peaks_location; None otherwise
        """
        real_idx = self.subset_idx[idx]
        ecg_raw = self.ecg_data[real_idx]      # shape e.g. [C, T] or [T, C]
        task_lbl = self.labels[idx]            # shape [num_classes]
        disease_lbl = self.disease_labels[idx] # shape [num_diseases]

        if self.use_peaks_location and self.rpeaks_data is not None:
            rpeaks_raw = self.rpeaks_data[real_idx]
            ecg_dict = {"ecg": ecg_raw, "rpeaks": rpeaks_raw}
            ecg_data = self.ecg_augmentations(ecg_dict)
            
            ecg_data_tensor = ecg_data["ecg"]
            rpeaks_t = ecg_data["rpeaks"]
            
        else:
            ecg_data_tensor = self.ecg_augmentations(ecg_raw)
            rpeaks_t = None

        task_labels_t = torch.from_numpy(task_lbl).float()
        disease_labels_t = torch.from_numpy(disease_lbl).float()

        return ecg_data_tensor, task_labels_t, disease_labels_t, rpeaks_t